#!/usr/bin/env python3
'''
把numpy卸载后可以正常运行，可能是numpy的版本问题吧

使用bellman的值迭代算法进行状态值的更新

bellman 更新的值迭代就是从当前状态，尝试所有的动作，每个动作所有能所达到的状态获得的激励+当前V值 总和 进行对比，最大的值更新当前状态的V值
'''
import gymnasium as gym
import collections # collections是用来替换python自带的list、dict等容器类的框架，提供了更多的功能，比如记录每个元素的访问次数
from tensorboardX import SummaryWriter

ENV_NAME = "FrozenLake-v1"
GAMMA = 0.9
TEST_EPISODES = 20


class Agent:
    def __init__(self):
        self.env = gym.make(ENV_NAME)
        self.state = self.env.reset()
        self.rewards = collections.defaultdict(float) # 奖励 key：(执行动作之前的状态、执行的动作、执行动作之后的状态) value: 获取的奖励
        self.transits = collections.defaultdict(collections.Counter) # 转换 key:(执行动作之前的状态、执行的动作) value: 执行动作后转换的状态 count: 该状态转换执行动作的次数
        self.values = collections.defaultdict(float) # key: 当前的状态 value: 所有动作执行后中最大的奖励，在值迭代中，这个就是V值

    def play_n_random_steps(self, count):
        '''
        开始随机探索

        count 随机探索的步数
        '''
        for _ in range(count):
            action = self.env.action_space.sample() # 从action_space随机选择一个动作
            new_state, reward, is_done, _ = self.env.step(action) # 通过环境执行一个动作 并获取新的状态、奖励、和是否已经结束
            self.rewards[(self.state, action, new_state)] = reward # 统计奖励信息
            self.transits[(self.state, action)][new_state] += 1 # 统计对应状态转换的次数
            self.state = self.env.reset() if is_done else new_state # 根据是否已经结束重置状态

    def calc_action_value(self, state, action):
        '''
        计算奖励累计值
        这里使用的是值迭代算法，计算 当前状态+执行的动作 => 目标状态后
        所有目标状态的所有奖励累计值

        state： 要计算累计值时的起始状态
        action: 状态所执行的动作
        '''
        # 获取当前 源状态+执行动作 = 对应状态 转换的次数
        # 因为转换到的目标状态可能不唯一，所以存在多个目标状态
        target_counts = self.transits[(state, action)]
        total = sum(target_counts.values())

        action_value = 0.0
        # 遍历每一个目标状态以及目标状态所转换的次数
        for tgt_state, count in target_counts.items():
            # 获取当前目标状态的奖励
            reward = self.rewards[(state, action, tgt_state)]
            # 计算动作值
            ## action_value: 在state状态执行action动作的奖励累计
            ## count / total: 转换到当前tgt_state的概率
            ## reward：奖励
            ## self.Value: 表示到达目标状态后，在目标状态可以转换为下一个状态中所能获得的最大奖励值
            action_value += (count / total) * (reward + GAMMA * self.values[tgt_state])
        # 返回动作值
        return action_value

    def select_action(self, state):
        '''
        根据当前状态
        从累计的奖励中，选择最大奖励的动作

        '''

        # best_action: 奖励最大的动作
        # best_value: 最大的奖励
        best_action, best_value = None, None
        # 遍历所有的动作空间
        for action in range(self.env.action_space.n):
            # 从累计的值中，获取当前的动作累计的奖励
            # 并判断当前的动作是否是最大的奖励并记录
            action_value = self.calc_action_value(state, action)
            if best_value is None or best_value < action_value:
                best_value = action_value
                best_action = action
        # 返回最大奖励的动作
        return best_action

    def play_episode(self, env):
        '''
        开启新一轮测试集，直到游戏结束退出返回

        返回本次测试获取的奖励累计值
        '''

        # 统计新一轮测试中的累计奖励
        total_reward = 0.0
        # reset 随机选择一个状态
        state = env.reset()
        while True:
            # 从累计的状态奖励表中获取最大奖励的动作
            action = self.select_action(state)
            # 执行该动作
            new_state, reward, is_done, _ = env.step(action)
            # 更新记录
            self.rewards[(state, action, new_state)] = reward
            self.transits[(state, action)][new_state] += 1
            total_reward += reward
            if is_done:
                break
            # 将新状态设置为当前状态
            state = new_state
        return total_reward

    def value_iteration(self):
        '''
        值遍历
        '''

        # self.env.observation_space.n 观察空间的维度
        # self.env.action_space.n 动作空间的维度
        # 遍历每一种观察空间的状态以及执行的动作计算每种观察空间和动作空间组合的累积值
        # 这里可以使用range是因为存储的dict可以直接用这种方式遍历
        for state in range(self.env.observation_space.n):
            # 使用的Bellman更新
            ## 在当前状态下遍历所有的动作，计算器转换的状态值
            ## 统计当前状态下执行所有可能的动作后所能到达的的目标状态后最大的可能奖励值
            ## 并存储起来
            state_values = [self.calc_action_value(state, action)
                            for action in range(self.env.action_space.n)]
            self.values[state] = max(state_values)


if __name__ == "__main__":
    test_env = gym.make(ENV_NAME)
    agent = Agent()
    writer = SummaryWriter(comment="-v-iteration")

    iter_no = 0
    best_reward = 0.0
    while True:
        iter_no += 1 # 遍历轮数
        agent.play_n_random_steps(100) # 随机探索
        agent.value_iteration() # 进行值遍历，更新每一个状态的最大奖励

        reward = 0.0
        # 开启新一轮测试，统计新一轮测试后的奖励累计值
        for _ in range(TEST_EPISODES):
            reward += agent.play_episode(test_env)
        # 计算平均每局奖励
        reward /= TEST_EPISODES
        writer.add_scalar("reward", reward, iter_no)
        # 如果当前测试获取的奖励大于已记录的最大奖励
        # 则打印并记录
        if reward > best_reward:
            print("Best reward updated %.3f -> %.3f" % (best_reward, reward))
            best_reward = reward
        if reward > 0.80:
            print("Solved in %d iterations!" % iter_no)
            break

    # 结束
    writer.close()
